import os
import shutil
import sys

import torch
from torch import nn
from torch import optim
from tqdm import tqdm

from config import Config as C
from model import resnet50
from utils import get_data, get_transform


def train(net, optimizer, loss_function, device, train_bar):
    # train
    net.train()
    r_loss = 0.0
    for step, data in enumerate(train_bar):
        images, labels = data
        optimizer.zero_grad()
        logits = net(images.to(device))
        loss = loss_function(logits, labels.to(device))
        loss.backward()
        optimizer.step()
        r_loss += loss.item()
        train_bar.desc = "loss:{:.3f}".format(loss)
    return r_loss


def val(net, device, val_bar):
    net.eval()
    acc = 0.0  # accumulate accurate number / epoch
    with torch.no_grad():
        for val_data in val_bar:
            val_images, val_labels = val_data
            outputs = net(val_images.to(device))
            # loss = loss_function(outputs, test_labels)
            predict_y = torch.max(outputs, dim=1)[1]
            acc += torch.eq(predict_y, val_labels.to(device)).sum().item()
    return acc


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    trans = get_transform()
    train_set, val_set, _, train_loader, val_loader, _ = get_data(
        C.ROOT_PATH,
        C.BATCH_SIZE,
        transform=trans,
        # target_transform=get_target_trans()
    )
    # train_set, test_set, train_loader, test_loader = get_data(
    #     C.ROOT_PATH,
    #     C.BATCH_SIZE,
    #     transform=trans,
    #     # target_transform=get_target_trans()
    # )
    val_num = len(val_set)
    # val_num = len(test_set)
    net = resnet50(num_classes=C.CLASSES_NUM)
    if os.path.exists(C.MODEL_WEIGHT_PATH):
        print('Find a checkpoint file, load state dict')
        shutil.copy2(C.MODEL_WEIGHT_PATH, os.path.join(
            os.path.dirname(C.MODEL_WEIGHT_PATH),
            os.path.basename(C.MODEL_WEIGHT_PATH) + '.old'
        ))
        net.load_state_dict(torch.load(C.MODEL_WEIGHT_PATH, map_location='cpu'))
    net.to(device)
    loss_function = nn.CrossEntropyLoss()
    params = [p for p in net.parameters() if p.requires_grad]
    optimizer = optim.Adam(params, lr=0.001)

    epochs = 30
    best_acc = 0.0
    save_path = C.MODEL_WEIGHT_PATH
    train_steps = len(train_loader)

    print(net)
    for epoch in range(epochs):
        train_bar = tqdm(train_loader, file=sys.stdout)
        # val_bar = tqdm(val_loader, file=sys.stdout)
        running_loss = train(net, optimizer, loss_function, device, train_bar)

        val_bar = tqdm(val_loader, file=sys.stdout)
        acc = val(net, device, val_bar)
        val_accurate = acc / val_num

        print('[epoch %d] train_loss: %.3f  val_accuracy: %.3f' % (epoch + 1, running_loss / train_steps, val_accurate))

        if val_accurate > best_acc:
            best_acc = val_accurate
            torch.save(net.state_dict(), save_path)


if __name__ == '__main__':
    main()
